# 并查集解决连接问题
class UnionFind:
    """
    每一个元素对应数组self.id的索引号
    每一个元素p所属的组为self.id[p]
    """
    def __init__(self, n):
        # 数组中元素的个数
        self.count = n
        # 使用数组的下标表示元素
        self.id = [0] * n
        for i in range(n):
            self.id[i] = i   # 每一个元素都是一个单独的组
    
    # 找到p所属的组
    # O(1) quick find
    def find(self, p):
        assert p >= 0 and p < self.count
        return self.id[p]

    # p与q是否相连
    def is_connected(self, p, q):
        return self.find(p) == self.find(q)

    # 合并两个元素为同一个组
    # 把p或者q所在的组的所有元素的ID号变成另一个元素的ID号
    # O(n)
    def union(self, p, q):
        id_p = self.find(p)
        id_q = self.find(q)
        # 遍历所有的元素把ID为id_p的所有元素的ID变成id_q
        for i in range(self.count):
            if self.id[i] == id_p:
                self.id[i] = id_q
